import os
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor

from utils import printer_util


class PrintJobDistributor:
    """打印任务分发器"""

    def __init__(self, config_manager, logger):
        self.config_manager = config_manager
        self.logger = logger
        # 线程池，最大并发数为5
        self.thread_pool = ThreadPoolExecutor(max_workers=5)
        # 用于跟踪正在执行的任务
        self.active_tasks = {}

    def distribute_jobs(self, files_to_print, update_ui_callback):
        """分发打印任务"""
        managed_printers = list(self.config_manager.get_printers().keys())
        if not managed_printers:
            return False, "没有配置任何打印机，无法执行打印任务。"

        # 只处理未打印的文件
        unprinted_files = [f for f in files_to_print if not f["printed"]]
        if not unprinted_files:
            return False, "没有需要打印的文件，所有文件都已打印。"

        # 清空之前的队列
        printer_queues = defaultdict(list)

        # 按文件类型分组
        file_groups = defaultdict(list)
        for file_info in unprinted_files:
            file_path = file_info["path"]
            ext = os.path.splitext(file_path)[1].lower()
            file_groups[ext].append(file_info)

        # 将文件按类型分配给打印机队列
        for ext, files in file_groups.items():
            # 为每种文件类型平均分配到各打印机
            files_per_printer = len(files) // len(managed_printers)
            remainder = len(files) % len(managed_printers)

            start_idx = 0
            for i, printer_name in enumerate(managed_printers):
                # 前remainder台打印机多分配一个文件
                end_idx = start_idx + files_per_printer + (1 if i < remainder else 0)
                if end_idx > start_idx:
                    printer_queues[printer_name].extend(files[start_idx:end_idx])
                    start_idx = end_idx

        # 为每个打印机队列启动处理任务
        for printer_name, files in printer_queues.items():
            # 跳过空队列
            if not files:
                continue

            # 按文件类型分组，以便批量处理
            printer_file_groups = defaultdict(list)
            for file_info in files:
                file_path = file_info["path"]
                ext = os.path.splitext(file_path)[1].lower()
                printer_file_groups[ext].append(file_info)

            # 为每种文件类型创建一个任务
            for ext, file_list in printer_file_groups.items():
                task_id = f"{printer_name}_{ext}"
                future = self.thread_pool.submit(
                    self.process_print_batch,
                    printer_name,
                    ext,
                    file_list,
                    update_ui_callback,
                )
                self.active_tasks[task_id] = future
                # 添加完成回调
                future.add_done_callback(
                    lambda f, tid=task_id: self.on_task_complete(f, tid)
                )

        return (
            True,
            f"{len(unprinted_files)} 个打印任务已分配到 {len(managed_printers)} 台打印机。",
        )

    def process_print_batch(
        self, printer_name, file_ext, file_list, update_ui_callback
    ):
        """处理同一类型的文件批量打印"""
        try:
            self.logger.info(
                f"开始批量处理 {len(file_list)} 个 {file_ext} 文件到打印机 {printer_name}"
            )

            # 获取打印机配置
            printer_config = self.config_manager.get_printer(printer_name)
            printer_ip = printer_config.get("ip", "")

            # 批量处理同类型文件
            for file_info in file_list:
                file_path = file_info["path"]
                # 根据文件类型获取对应的打印配置
                file_config = self.config_manager.get_config_for_file(file_path)
                # 合并打印机IP和文件配置
                print_params = {"ip": printer_ip, **file_config}

                # 执行打印
                printer_util.print_file(printer_name, file_path, print_params)

                # 标记文件为已打印
                if update_ui_callback:
                    update_ui_callback(file_info["id"])

            self.logger.info(
                f"完成批量处理 {len(file_list)} 个 {file_ext} 文件到打印机 {printer_name}"
            )
            return True

        except Exception as e:
            self.logger.error(f"批量打印失败: {e}")
            return False

    def on_task_complete(self, future, task_id):
        """任务完成回调"""
        try:
            # 获取任务结果
            result = future.result()
            if result:
                self.logger.info(f"任务 {task_id} 完成")
            else:
                self.logger.warning(f"任务 {task_id} 失败")
        except Exception as e:
            self.logger.error(f"任务 {task_id} 执行异常: {e}")
        finally:
            # 从活动任务中移除
            if task_id in self.active_tasks:
                del self.active_tasks[task_id]

    def shutdown(self):
        """关闭线程池"""
        # 取消所有活动任务
        for task_id, future in self.active_tasks.items():
            future.cancel()
        # 关闭线程池
        self.thread_pool.shutdown(wait=False)
